import time, pickle

import numpy as np
from paddleocr import PaddleOCR


class ScreenReader():
    def __init__(self):
        self.roi = (1990, 312, 2400, 734)
        self.sub_roi_pos = np.asarray([
            [26, 9, 218, 37],  # 套装
            [23, 58, 79, 79],  # 位置
            [22, 123, 78, 148],  # 主词条
            [25, 258, 73, 280],  # 等级
            [30, 301, 300, 322],  # 副词条1
            [30, 334, 300, 354],  # 副词条2
            [30, 365, 300, 385],  # 副词条3
            [30, 397, 300, 418]  # 副词条4
        ])
        self.lock_pos = [275, 268]
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch", show_log=False, use_gpu=False)
        pass

    def read_artifact(self, pillow_img):
        start = time.time()
        roi = pillow_img.crop(self.roi)
        np_array = np.asarray(roi)
        lock_color = np_array[self.lock_pos[0], self.lock_pos[1], 0:3]
        if np.linalg.norm(lock_color - [255, 137, 117]) <= 10:
            is_locked = True
        else:
            is_locked = False
        str_list = self.ocr.ocr(np_array, cls=False)[0]
        if str_list == None:
            return None
        str_pos_list = []
        for str in str_list:
            box = np.asarray(str[0])
            x_min = np.min(box[:, 0])
            x_max = np.max(box[:, 0])
            y_min = np.min(box[:, 1])
            y_max = np.max(box[:, 1])
            local_pos = np.asarray([x_min, y_min, x_max, y_max])
            str_pos_list.append(local_pos)
        sub_roi_num = self.sub_roi_pos.shape[0]
        str_num = len(str_pos_list)
        fit_matrix = np.zeros((sub_roi_num, str_num))
        roi_str_list = []
        for roi_idx in range(sub_roi_num):
            for str_idx in range(str_num):
                fit_x_min = np.max([self.sub_roi_pos[roi_idx][0], str_pos_list[str_idx][0]])
                fit_y_min = np.max([self.sub_roi_pos[roi_idx][1], str_pos_list[str_idx][1]])
                fit_x_max = np.min([self.sub_roi_pos[roi_idx][2], str_pos_list[str_idx][2]])
                fit_y_max = np.min([self.sub_roi_pos[roi_idx][3], str_pos_list[str_idx][3]])
                fit_x_size = np.clip(fit_x_max - fit_x_min, 0, None)
                fit_y_size = np.clip(fit_y_max - fit_y_min, 0, None)
                fit_matrix[roi_idx, str_idx] = fit_x_size * fit_y_size
            match_str_idx = np.argsort(fit_matrix[roi_idx, :])[-1]
            if fit_matrix[roi_idx, match_str_idx] >= 200:
                roi_str_list.append(str_list[match_str_idx][1][0])
        if len(roi_str_list) < 5 or len(roi_str_list) > 8:
            return None
        set_name = roi_str_list[0]
        set_pos = roi_str_list[1]
        main_affix = roi_str_list[2]
        level = int(roi_str_list[3])
        aux_affix_list = []
        for str in roi_str_list[4:]:
            if str[0] == '·':
                str = str[1:]
                local_affix = str.split('+', 1)
                if local_affix[1][-1] == '%':
                    local_affix[1] = 0.01 * float(local_affix[1][:-1])
                else:
                    local_affix[1] = local_affix[1].replace("，", "")
                    local_affix[1] = local_affix[1].replace(",", "")
                    local_affix[1] = float(local_affix[1])
                aux_affix_list.append(local_affix)
        print('ocr cost time %.3f' % (time.time() - start))
        return [set_name, set_pos, main_affix, level, is_locked, aux_affix_list]

    def get_artifact_score(self, artifact):
        aux_affix_score = 0.0
        for aux_affix in artifact[5]:
            if aux_affix[0] == '暴击率':
                aux_affix_score += aux_affix[1] / 0.01
            elif aux_affix[0] == '暴击伤害':
                aux_affix_score += aux_affix[1] / 0.02
            elif aux_affix[0] == '元素充能效率':
                aux_affix_score += aux_affix[1] / 0.0167
            elif aux_affix[0] == '攻击力':
                if aux_affix[1] < 1:
                    aux_affix_score += aux_affix[1] / 0.015
            elif aux_affix[0] == '元素精通':
                aux_affix_score += aux_affix[1] / 6
            elif aux_affix[0] == '生命值':
                if aux_affix[1] < 1:
                    aux_affix_score += aux_affix[1] / 0.015
            elif aux_affix[0] == '防御力':
                if aux_affix[1] < 1:
                    aux_affix_score += aux_affix[1] / 0.0187
            else:
                raise Exception('副词条识别错误')
        return aux_affix_score

    def get_effective_affix_num(self, artifact):
        effective_aux_affix_num = 0
        for aux_affix in artifact[5]:
            if aux_affix[0] == '暴击率':
                effective_aux_affix_num += 1
            elif aux_affix[0] == '暴击伤害':
                effective_aux_affix_num += 1
            elif aux_affix[0] == '元素充能效率':
                effective_aux_affix_num += 0.5
            elif aux_affix[0] == '攻击力':
                if aux_affix[1] < 1:
                    effective_aux_affix_num += 0.5
            elif aux_affix[0] == '元素精通':
                effective_aux_affix_num += 0.5
            elif aux_affix[0] == '生命值':
                if aux_affix[1] < 1:
                    effective_aux_affix_num += 0.5
            elif aux_affix[0] == '防御力':
                if aux_affix[1] < 1:
                    effective_aux_affix_num += 0.5
            else:
                raise Exception('副词条识别错误')
        return effective_aux_affix_num

    def artifact_filter(self, artifact):
        main_affix = artifact[2]
        key_affix_list = ['暴击', '充能', '精通', '伤害加成', '治疗']
        is_keep = False
        for key_affix in key_affix_list:
            is_keep = is_keep or (key_affix in main_affix)
        if is_keep:
            return is_keep, None, None
        level = artifact[3]
        if level == 20:
            score = self.get_artifact_score(artifact)
            if score >= 20:
                is_keep = True
            else:
                is_keep = False
            return is_keep, score, None
        else:
            effective_aux_affix_num = self.get_effective_affix_num(artifact)
            aux_affix_num = len(artifact[5])
            if aux_affix_num == 4 and effective_aux_affix_num >= 2:
                is_keep = True
            elif aux_affix_num == 3 and effective_aux_affix_num >= 1.5:
                is_keep = True
            else:
                is_keep = False
            return is_keep, None, effective_aux_affix_num

# screen_reader = ScreenReader()
# artifact_list = []
# for i in range(0, 1628):
#     filename = 'artifact_img\\artifact_%04d.jpg' % (i)
#     img = Image.open(filename)
#     artifact_list.append(screen_reader.read_artifact(img))
#     print('artifact %d / %d' % (i, 1627))
# print(artifact_list)
# with open('artifact_list.pickle', 'wb') as file:
#     # 使用pickle.dump()函数将列表序列化并保存到文件
#     pickle.dump(artifact_list, file)

# with open('artifact_list.pickle', 'rb') as file:
#     # 使用pickle.load()函数从文件中读取并反序列化列表
#     artifact_list = pickle.load(file)
#
# keep_list=[]
# score_list = []
# affix_num_list=[]
# for artifact in artifact_list:
#     is_keep,score,affix_num = screen_reader.artifact_filter(artifact)
#     keep_list.append(is_keep)
#     if score != None:
#         score_list.append(score)
#     if affix_num != None:
#         affix_num_list.append(affix_num)
# keep_list=np.asarray(keep_list)
# score_list = np.asarray(score_list)
# affix_num_list=np.asarray(affix_num_list)
# print(np.sum(keep_list))
# print(score_list)
# print(affix_num_list)
